import astropy.io.fits as fits
import math
import numpy as np
from astropy.wcs import WCS
from astropy.coordinates import SkyCoord
import astropy.units as u
import pickle
from scipy import interpolate
import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib.ticker import (MultipleLocator, AutoMinorLocator)
c=3.0e10

def rms(x, axis=None):
    return np.sqrt(np.mean(x**2, axis=axis))


spectral_template_dict=pickle.load(open("spectral_dict_lime_stacked.pickle", "rb"))

v_src=-0.0
template_trans={
                'line1' : {'freq' : 203.40752e9, 'A_ul': 4.812E-06,'E_u': 203.7, 'g_u':7 , 'trans': ''}
               }

for key in spectral_template_dict.keys():
   rel_freq=spectral_template_dict[key]['freq_axis']*1e9-template_trans['line1']['freq']
   spectral_template_dict[key]['rel_freq_axis']=spectral_template_dict[key]['freq_axis']*1e9-template_trans['line1']['freq']


   fig,ax=plt.subplots(1,1,figsize=(15, 6))
   ax.plot(spectral_template_dict[key]['rel_freq_axis']/1e9,spectral_template_dict[key]['spectrum'],drawstyle='steps',linewidth='1')
   plt.savefig(key+'_spectrum.png')
   spectral_template_dict[key]['spectrum']=spectral_template_dict[key]['spectrum']/np.max(spectral_template_dict[key]['spectrum'])

   new_avg_spectrum_wflip=(spectral_template_dict[key]['spectrum']+np.flip(spectral_template_dict[key]['spectrum']))/2.0

   new_avg_spectrum_wflip=new_avg_spectrum_wflip/np.max(new_avg_spectrum_wflip)

   fig,ax=plt.subplots(1,1,figsize=(15, 6))
   ax.plot(spectral_template_dict[key]['rel_freq_axis']/1e9,spectral_template_dict[key]['spectrum'],drawstyle='steps',linewidth='1')
   ax.plot(spectral_template_dict[key]['rel_freq_axis']/1e9,np.flip(spectral_template_dict[key]['spectrum']),drawstyle='steps',linewidth='1')
   ax.plot(spectral_template_dict[key]['rel_freq_axis']/1e9,new_avg_spectrum_wflip,drawstyle='steps',linewidth='1')
   plt.savefig('flipped_'+key+'_spectra.png')


spectral_template_lime={'freq_axis': spectral_template_dict[key]['rel_freq_axis'].copy(),'spectrum': new_avg_spectrum_wflip.copy() }


pickle.dump(spectral_template_lime, open("spectral_template_lime_stacked.pickle", "wb"))







